// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/dwio/common/column_selector.h>

#include <iterator>

#include "melon/conv.h"
#include <pollux/dwio/common/type_utils.h>

namespace kumo::pollux::dwio::common {
    using pollux::RowType;
    using pollux::Type;
    using pollux::TypeKind;

    namespace {
        std::string childName(
            const std::shared_ptr<const pollux::Type> &type,
            size_t childIdx,
            const std::string &inherit) {
            switch (type->kind()) {
                case pollux::TypeKind::ARRAY: {
                    DWIO_ENSURE_EQ(childIdx, 0);
                    return melon::to<std::string>(inherit, ".[ITEM]");
                }
                case pollux::TypeKind::MAP: {
                    DWIO_ENSURE_LT(childIdx, 2);
                    return melon::to<std::string>(
                        inherit, childIdx == 0 ? ".[KEY]" : ".[VALUE]");
                }
                case pollux::TypeKind::ROW: {
                    auto name = std::dynamic_pointer_cast<const pollux::RowType>(type)->nameOf(
                        childIdx);
                    return melon::to<std::string>(inherit, ".", name);
                }
                default:
                    // for others that has no sub field filter - just use its parent's name
                    return inherit;
            }
        }
    } // namespace

    void ColumnSelector::buildNodes(
        const std::shared_ptr<const RowType> &schema,
        const std::shared_ptr<const RowType> &contentSchema) {
        buildNode(
            FilterNode(0, MAX_UINT64, "_ROOT_", "", false),
            nullptr,
            std::dynamic_pointer_cast<const Type>(schema),
            std::dynamic_pointer_cast<const Type>(contentSchema),
            true);
    }

    FilterTypePtr ColumnSelector::buildNode(
        const FilterNode &node,
        const std::shared_ptr<FilterType> &parent,
        const std::shared_ptr<const Type> &type,
        const std::shared_ptr<const Type> &contentType,
        const bool inContent) {
        DWIO_ENSURE(node.node == nodes_.size(), "current node align with vector");

        // make sure content type is compatible to reading request type
        // when content type is present (it may be absent due to schema mismatch)
        if (contentType != nullptr) {
            typeutils::checkTypeCompatibility(*contentType, *type);
        }

        // the process will cover supported schema evolution:
        // appending columns at the end or deleting columns at the end
        auto current =
                std::make_shared<FilterType>(node, parent, inContent, type, contentType);
        nodes_.push_back(current);

        // generate nodes for all its children
        // check if this requested type is in the content to be read or not
        // also, if the asked type shorter than content, it won't be show up in
        // column selector filter tree
        nodes_.reserve(nodes_.size() + type->size());
        if (node.node == 0) {
            auto &rowType = type->as_row();
            for (size_t i = 0, size = type->size(); i < size; ++i) {
                bool inData = contentType && i < contentType->size();
                current->addChild(buildNode(
                    FilterNode(nodes_.size(), i, rowType.nameOf(i), "", !inData),
                    current,
                    type->childAt(i),
                    inData ? contentType->childAt(i) : nullptr,
                    inData));
            }
        } else {
            for (size_t i = 0, size = type->size(); i < size; ++i) {
                bool inData = contentType && i < contentType->size();
                current->addChild(buildNode(
                    FilterNode(
                        nodes_.size(),
                        node.column,
                        childName(type, i, node.name),
                        "",
                        !inData),
                    current,
                    type->childAt(i),
                    inData ? contentType->childAt(i) : nullptr,
                    inData));
            }
        }

        return current;
    }

    // this copy method only update inContent and data type
    // based on disk data type
    void ColumnSelector::copy(
        common::FilterTypePtr &node,
        const std::shared_ptr<const Type> &diskType,
        const common::FilterTypePtr &origin) {
        auto originIsNull = (origin == nullptr);
        if (!originIsNull) {
            node->setInContent(origin->isInContent());
            node->setSequenceFilter(origin->getSequenceFilter());
        }

        // we're not interested in non-read nodes and nodes that not in content
        if (node->shouldRead() && node->isInContent()) {
            // not found in disk type
            if (diskType == nullptr) {
                node->setInContent(false);
                return;
            }

            // ensure disk type can be converted to request type
            typeutils::checkTypeCompatibility(*diskType, *node->getRequestType());

            // update data type during the visit as well as other data fields
            node->setDataType(diskType);
            if (!originIsNull) {
                const common::FilterNode &f = origin->getNode();
                auto &fn = const_cast<common::FilterNode &>(node->getNode());
                fn.expression = f.expression;
                fn.partitionKey = f.partitionKey;
            }

            // visit all children
            for (size_t i = 0; i < node->size(); ++i) {
                copy(
                    const_cast<common::FilterTypePtr &>(node->childAt(i)),
                    i < diskType->size() ? diskType->childAt(i) : nullptr,
                    originIsNull ? nullptr : origin->childAt(i));
            }
        }
    }

    ColumnSelector ColumnSelector::apply(
        const std::shared_ptr<ColumnSelector> &origin,
        const std::shared_ptr<const RowType> &fileSchema) {
        // current instance maybe null.
        if (origin == nullptr) {
            return ColumnSelector(fileSchema);
        }

        // if selector has no schema, we just build a new tree with file schema
        // selector.getProjection will carry all logic information including nodes
        const bool onlyFilter = !origin->hasSchema();
        ColumnSelector cs(
            onlyFilter ? fileSchema : origin->getSchema(),
            origin->getNodeFilter(),
            true);

        // visit file schema and decide in content call
        copy(cs.nodes_[0], fileSchema, onlyFilter ? nullptr : origin->nodes_[0]);
        return cs;
    }

    /**
     * Mark the node and all its children recursively to be read.
     *
     * @param node the starting id
     */
    void ColumnSelector::setRead(const common::FilterTypePtr &node, bool only) {
        if (!node->valid()) {
            return;
        }

        // and all its child
        node->setRead();
        if (only) {
            return;
        }

        // set all its children to be in read state
        for (size_t i = 0, size = node->size(); i < size; ++i) {
            setRead(node->childAt(i));
        }

        // and all its ancestor should be in read state
        // so a reader is expected to check if reading a top node
        // if no, none of its children should be read
        // if yes, all or part of its children may be read - need further check on
        // node
        for (auto &&p = node->getParent().lock(); p; p = p->getParent().lock()) {
            p->setRead();
        }
    }

    // TODO (cao) - SelectedTypeBuilder is uncessary to exist, can be killed
    std::shared_ptr<const TypeWithId> ColumnSelector::buildSelected() const {
        auto selector = [this](size_t index) { return shouldReadNode(index); };
        return typeutils::buildSelectedType(TypeWithId::create(schema_), selector);
    }

    std::shared_ptr<const RowType> ColumnSelector::buildSelectedReordered() const {
        std::vector<std::string> names;
        std::vector<std::shared_ptr<const Type> > types;
        for (auto &node: filter_) {
            names.push_back(schema_->nameOf(node.column));
            types.push_back(schema_->childAt(node.column));
        }
        return std::make_shared<RowType>(std::move(names), std::move(types));
    }

    void ColumnSelector::setReadAll() {
        DWIO_ENSURE(shouldReadAll(), "should not called on non-read all case");
        auto root = getNode(0);
        setRead(root);

        // generate all columns in the filter list
        for (size_t i = 0, size = root->size(); i < size; ++i) {
            filter_.push_back(root->childAt(i)->getNode());
        }
    }

    const ColumnFilter &ColumnSelector::getProjection() const {
        return filter_;
    }

    void ColumnSelector::setConstants(
        const std::vector<std::string> &keys,
        const std::vector<std::string> &values) {
        DWIO_ENSURE_EQ(values.size(), keys.size(), "keys/values size mismatch");

        // for every node value pair - set the value into the node's expression
        for (size_t i = 0, size = keys.size(); i < size; ++i) {
            auto &key = keys[i];
            auto &value = values[i];

            auto node = findColumn(key);
            DWIO_ENSURE(node->valid(), "constant node not found: ", key);

            // currently we only support limited set of types for constant expression
            auto kind = node->getKind();
            DWIO_ENSURE(
                kind == TypeKind::BOOLEAN || kind == TypeKind::TINYINT ||
                kind == TypeKind::SMALLINT || kind == TypeKind::INTEGER ||
                kind == TypeKind::BIGINT || kind == TypeKind::REAL ||
                kind == TypeKind::DOUBLE || kind == TypeKind::VARCHAR,
                "unsupported constant type: ",
                kind);

            // assign the value to the constant's expression
            const_cast<FilterNode &>(node->getNode()).expression = value;
        }
    }

    // note that - this list includes root NODE - should we?
    std::vector<uint64_t> ColumnSelector::getNodeFilter() const {
        std::vector<uint64_t> nodeIds;
        nodeIds.reserve(nodes_.size());
        for (const auto &node: nodes_) {
            if (node->shouldRead()) {
                nodeIds.emplace_back(node->getId());
            }
        }
        return nodeIds;
    }

    const FilterTypePtr &ColumnSelector::process(const std::string &column, bool) {
        // map column support column name followed by expression
        // after "#" character which is not allowed as legal column name
        auto pair = extractColumnName(column);
        auto expr = pair.second;
        if (!expr.empty()) {
            // when seeing this type of column, we require it has to exist
            const auto &node = findNode(pair.first);
            if (node->valid()) {
                DWIO_ENSURE(
                    node->getKind() == TypeKind::MAP || node->getKind() == TypeKind::ROW,
                    "only support expression for map or row type currently: ",
                    column);

                // set expression for this node
                auto &nodeValue = node->getNode();
                const_cast<common::FilterNode &>(nodeValue).expression = expr;
            }
            return node;
        }

        // pass over special rules handling - find the node and return it
        return findNode(column);
    }

    std::pair<std::string_view, std::string_view> extractColumnName(
        const std::string_view &name) {
        // right now this is the only supported expression for MAP key filter
        const auto pos = name.find('#');
        if (pos == std::string::npos) {
            return std::make_pair(name, "");
        }

        // Special map column handling.
        const auto colName = name.substr(0, pos);
        const auto expr = name.substr(pos + 1);
        return std::make_pair(colName, expr);
    }

    void ColumnSelector::logFilter() const {
        auto numColumns = 0;
        auto numNodes = 0;
        // log current filter_ members
        if (hasSchema()) {
            numColumns = getNode(0)->size();
            numNodes = nodes_.size();
        }

        // log it out
        getLog()->logColumnFilter(filter_, numColumns, numNodes, hasSchema());
    }

    std::shared_ptr<ColumnSelector> ColumnSelector::fromScanSpec(
        const pollux::common::ScanSpec &spec,
        const RowTypePtr &rowType) {
        std::vector<std::string> columnNames;
        for (auto &child: spec.children()) {
            if (!child->readFromFile()) {
                continue;
            }
            std::string name = child->fieldName();
            if (!child->flatMapFeatureSelection().empty()) {
                const auto &featureIds = child->flatMapFeatureSelection();
                const auto &type = rowType->findChild(name);
                POLLUX_CHECK(type->is_map());
                const bool is_varchar = type->as_map().keyType()->is_varchar();
                name += "#[";
                for (int i = 0; i < featureIds.size(); ++i) {
                    if (i > 0) {
                        name += ',';
                    }
                    if (is_varchar) {
                        melon::json::escapeString(featureIds[i], name, {});
                    } else {
                        name += featureIds[i];
                    }
                }
                name += ']';
            }
            columnNames.push_back(std::move(name));
        }
        if (columnNames.empty()) {
            static const RowTypePtr kEmpty{ROW({}, {})};
            return std::make_shared<ColumnSelector>(kEmpty);
        }
        return std::make_shared<ColumnSelector>(rowType, columnNames);
    }
} // namespace kumo::pollux::dwio::common
